//
// Tabular Q-RewardMachine
//
// lijin@air.tsinghua.edu.cn
//

#include <cv.h>
#include <highgui.h>
#include <opencv2/opencv.hpp>
#include <cstdlib>

#include "Environment.hpp"

using namespace cv;
using namespace std;

int delay_ms = 1;

void DrawEnvironment(Environment env)
{
    cv::Mat map_( 1000, 1000, CV_8UC3, Scalar(255,255,255) );
    
    // 1. 绘制grid map
    double grid_sz = map_.cols/MAP_W;
    for(int x=0;x<MAP_W;x++)
        for (int y=0;y<MAP_H;y++)
        {
            if(env.map[x][y]==0)
                cv::rectangle(map_,cv::Point(x*grid_sz,y*grid_sz),cv::Point(x*grid_sz+grid_sz,y*grid_sz+grid_sz),Scalar(0,230,0));
            else if(env.map[x][y]==1)
            {
                cv::rectangle(map_,cv::Point(x*grid_sz,y*grid_sz),cv::Point(x*grid_sz+grid_sz,y*grid_sz+grid_sz),Scalar(0,230,0));
                cv::rectangle(map_,cv::Point(x*grid_sz,y*grid_sz),cv::Point(x*grid_sz+grid_sz,y*grid_sz+grid_sz),Scalar(100,100,100),-1);
            }
            else if(env.map[x][y]==2)
            {
                cv::rectangle(map_,cv::Point(x*grid_sz,y*grid_sz),cv::Point(x*grid_sz+grid_sz,y*grid_sz+grid_sz),Scalar(0,230,0));
                cv::rectangle(map_,cv::Point(x*grid_sz,y*grid_sz),cv::Point(x*grid_sz+grid_sz,y*grid_sz+grid_sz),Scalar(0,0,255),-1);
            }
        }
    double agent_x = env.smart_agent.state.x;
    double agent_y = env.smart_agent.state.y;
    cv::rectangle(map_,cv::Point(agent_x*grid_sz,agent_y*grid_sz),cv::Point(agent_x*grid_sz+grid_sz,agent_y*grid_sz+grid_sz),Scalar(0,230,0));
    cv::rectangle(map_,cv::Point(agent_x*grid_sz,agent_y*grid_sz),cv::Point(agent_x*grid_sz+grid_sz,agent_y*grid_sz+grid_sz),Scalar(255,0,0),-1);

    double start_x = env.start.x;
    double start_y = env.start.y;
    cv::rectangle(map_,cv::Point(start_x*grid_sz,start_y*grid_sz),cv::Point(start_x*grid_sz+grid_sz,start_y*grid_sz+grid_sz),Scalar(255,0,0),3);

    double wp_x = env.WP.x;
    double wp_y = env.WP.y;
    cv::rectangle(map_,cv::Point(wp_x*grid_sz,wp_y*grid_sz),cv::Point(wp_x*grid_sz+grid_sz,wp_y*grid_sz+grid_sz),Scalar(0,230,0),3);


    // 2. 绘制Agent策略
    cv::Scalar color_of_policy(230,0,0);
    for(int x=0;x<MAP_W;x++)
        for (int y=0;y<MAP_H;y++)
        {
            if(env.map[x][y]!=0)continue;
            // 使用四个顶点指示action分布
            cv::Point2f up,left,down,right,center;
            center = cv::Point2f(x*grid_sz+0.5*grid_sz,y*grid_sz+0.5*grid_sz);
            up = left = down = right = center;
            double pi[4] = {env.smart_agent.Policy_start_to_waypoint[x][y].pi[0],
                            env.smart_agent.Policy_start_to_waypoint[x][y].pi[1],
                            env.smart_agent.Policy_start_to_waypoint[x][y].pi[2],
                            env.smart_agent.Policy_start_to_waypoint[x][y].pi[3]};
            double l = sqrt(pi[0]*pi[0]+pi[1]*pi[1]+pi[2]*pi[2]+pi[3]*pi[3]);
            pi[0] /= l ;pi[0] *= (0.5*grid_sz);
            pi[1] /= l ;pi[1] *= (0.5*grid_sz);
            pi[2] /= l ;pi[2] *= (0.5*grid_sz);
            pi[3] /= l ;pi[3] *= (0.5*grid_sz);

            pi[0] = max(0.0,pi[0]);
            pi[1] = max(0.0,pi[1]);
            pi[2] = max(0.0,pi[2]);
            pi[3] = max(0.0,pi[3]);

            up.y -= pi[0];
            left.x -= pi[1];
            down.y += pi[2];
            right.x += pi[3];
            
            cv::line(map_,up,left,color_of_policy,1,CV_AA);
            cv::line(map_,left,down,color_of_policy,1,CV_AA);
            cv::line(map_,down,right,color_of_policy,1,CV_AA);
            cv::line(map_,right,up,color_of_policy,1,CV_AA);
        }

    // 2. 绘制Agent策略
    color_of_policy = cv::Scalar(0,0,230);
    for(int x=0;x<MAP_W;x++)
        for (int y=0;y<MAP_H;y++)
        {
            if(env.map[x][y]!=0)continue;
            // 使用四个顶点指示action分布
            cv::Point2f up,left,down,right,center;
            center = cv::Point2f(x*grid_sz+0.5*grid_sz,y*grid_sz+0.5*grid_sz);
            up = left = down = right = center;
            double pi[4] = {env.smart_agent.Policy_waypoint_to_goal[x][y].pi[0],
                            env.smart_agent.Policy_waypoint_to_goal[x][y].pi[1],
                            env.smart_agent.Policy_waypoint_to_goal[x][y].pi[2],
                            env.smart_agent.Policy_waypoint_to_goal[x][y].pi[3]};
            double l = sqrt(pi[0]*pi[0]+pi[1]*pi[1]+pi[2]*pi[2]+pi[3]*pi[3]);
            pi[0] /= l ;pi[0] *= (0.5*grid_sz);
            pi[1] /= l ;pi[1] *= (0.5*grid_sz);
            pi[2] /= l ;pi[2] *= (0.5*grid_sz);
            pi[3] /= l ;pi[3] *= (0.5*grid_sz);

            pi[0] = max(0.0,pi[0]);
            pi[1] = max(0.0,pi[1]);
            pi[2] = max(0.0,pi[2]);
            pi[3] = max(0.0,pi[3]);

            up.y -= pi[0];
            left.x -= pi[1];
            down.y += pi[2];
            right.x += pi[3];
            
            cv::line(map_,up,left,color_of_policy,1,CV_AA);
            cv::line(map_,left,down,color_of_policy,1,CV_AA);
            cv::line(map_,down,right,color_of_policy,1,CV_AA);
            cv::line(map_,right,up,color_of_policy,1,CV_AA);
        }

    std::string text = std::to_string(env.smart_agent.success_rate);
	int font_face = cv::FONT_HERSHEY_COMPLEX; 
	double font_scale = 2;
	int thickness = 2;
	int baseline;
	//获取文本框的长宽
	cv::Size text_size = cv::getTextSize(text, font_face, font_scale, thickness, &baseline);
    cv::Point origin; 
	origin.x = 100;
	origin.y = 100;
	cv::putText(map_, text, origin, font_face, font_scale, cv::Scalar(0, 0, 0), thickness, 8);

    cv::imshow("map_",map_);
    if(cv::waitKey(delay_ms)==(int)'q')exit(0);
    map_.release();
}

int main(int argc, char* argv[])
{
    namedWindow("map_");

    Environment Env;

    int update_counter = 0;

    for(;;)
    {
        Env.step();

        if(Env.smart_agent.success_rate >= 0.9)
        {
            // slow down for inspection
            delay_ms = 20;
            Env.smart_agent.epsilon = 0;

            std::cout<<"update_counter : "<<update_counter<<std::endl;
        }
        else update_counter++;

        DrawEnvironment(Env);
    }

    return 0;
}
